#!/usr/bin/python
import os
import random
import shutil

def SplitTrainTestDataSet(src_dir):
    if not os.path.exists(os.path.join(src_dir, "train")):
        os.makedirs(os.path.join(src_dir, "train"))
    if not os.path.exists(os.path.join(src_dir, "test")):
        os.makedirs(os.path.join(src_dir, "test"))
    train_test_dir_depth = len(os.path.join(src_dir, "train").strip().split("/")) #strip删除字符，当括号为空时，默认删除空白符
    for root, dirs, files in os.walk(src_dir):
        root_split = root.strip().split("/")
        label_name =  root_split[-1]
        if label_name == src_dir.strip().split("/")[-1] \
            or (train_test_dir_depth == len(root_split) and "train" == label_name) \
            or (train_test_dir_depth == len(root_split) and "test" == label_name) \
            or train_test_dir_depth < len(root_split):
            continue

        train_data_file = os.path.join(src_dir, "train/" + label_name)
        test_data_file = os.path.join(src_dir, "test/" + label_name)
        if not os.path.exists(train_data_file):
            os.makedirs(train_data_file)
        if not os.path.exists(test_data_file):
            os.makedirs(test_data_file)

        random.shuffle(files)
        # print(len(files))
        file_size = len(files)
        for i in range(file_size):
            if i < int(0.9*file_size):
                shutil.move(os.path.join(root, files[i]), train_data_file)
            else:
                shutil.move(os.path.join(root, files[i]), test_data_file)
        os.removedirs(root)
    return

if __name__ == "__main__":
    src_dir = "/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/data/"#根目录
    SplitTrainTestDataSet(src_dir)
        

